package org.boot.chat.core;

import lombok.extern.slf4j.Slf4j;
import org.boot.chat.config.ChatModelsProperties;
import org.boot.chat.config.EasyChatModelsFactory;
import org.boot.chat.data.model.entity.ChatModelDO;
import org.boot.chat.util.SpringUtil;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import java.util.Arrays;

/**
 * EasyChat核心操作类
 * @author 最爱芫荽
 * @email 15733440205@163.com
 */
@Slf4j
public class EasyChatUtil {

    private static ChatModelsProperties chatModelsProperties;

    private static EasyChatModelsFactory chatModelsFactory;

    static {
        chatModelsProperties = SpringUtil.getBean(ChatModelsProperties.class);
        chatModelsFactory = SpringUtil.getBean(EasyChatModelsFactory.class);
    }

    public static String call(String model, String message) {
        Prompt prompt = new Prompt(new UserMessage(message));
        Generation generation = chatModel(model).call(prompt).getResult();
        return generation != null ? generation.getOutput().getText() : "";
    }

    public static String call(String model, Message... messages) {
        Prompt prompt = new Prompt(Arrays.asList(messages));
        Generation generation = chatModel(model).call(prompt).getResult();
        return generation != null ? generation.getOutput().getText() : "";
    }


    // * - 获取OpenAiChatModel实例 - *//
    public static OpenAiChatModel chatModel(String model) {
        if (model != null){
            String beanName = getBeanName(model);
            if (beanName != null){
                log.info("use model: {}", beanName);
                return SpringUtil.getBean(beanName, OpenAiChatModel.class);
            }
        }
        log.info("model is null, use default model: {}", getDefaultBeanName());
        return SpringUtil.getBean(getDefaultBeanName(), OpenAiChatModel.class);
    }

    // * - 重载模型列表 - *//
    public void overloadChatModel(){
        chatModelsFactory.initBean();
    }

    //* - 根据模型名称获取bean名称 - *//
    public static String getBeanName(String model){
        for (ChatModelDO chatModel : chatModelsProperties.getChatModels()) {
            String bean = chatModel.getBean();
            if (bean.equals(model)){
                return bean;
            }
        }
        return null;
    }

    //* - 获取默认模型的bean名称 - *//
    public static String getDefaultBeanName(){
        for (ChatModelDO chatModel : chatModelsProperties.getChatModels()) {
            if (chatModel.getPrimary().equals(0)){
                return chatModel.getBean();
            }
        }
        return null;
    }
}
